import { RRDBNet } from '../RRDBNet'
import tf from '../tf'
import { BaseConvolutionModel } from './BaseConvolutionModel'
import BaseLoadHandler from './loadHandler/BaseLoadHandler'
import SimpleLoadHandler from './loadHandler/SimpleLoadHandler'
import { upscaleModelUrls } from './urls'

function loadModel(weights: any) {
  const rRDBNet = new RRDBNet({
    numInCh: 3,
    numOutCh: 3,
    scale: 4,
    numBlock: 23,
  })
  const _input = tf.input({ shape: [null, null, 3] })
  const model = tf.model({
    inputs: _input,
    outputs: rRDBNet.apply(_input) as any,
  })
  // const weights = await loadMsgpackJson('/model/RealESRGAN_x4plus.bin')
  const tensorWeights = weights.map((w: any[]) => {
    if (Array.isArray(w[0])) {
      // weight
      return tf.tidy(() => {
        return tf.tensor(w).transpose([2, 3, 1, 0])
      })
    } else {
      //bias
      return tf.tensor(w, [w.length])
    }
  })
  // console.log(weights)
  model.setWeights(tensorWeights)
  return model
}
class RealESRGANModel extends BaseConvolutionModel {
  private loadHander: BaseLoadHandler
  constructor() {
    super()
    this.loadHander = new SimpleLoadHandler(
      'RealESRGANModel',
      upscaleModelUrls.RealESRGAN_x4plus,
      1024 * 1024 * 150,
      async (data) => {
        this.model = loadModel(data)
      }
    )
    // this.loadHander
    //   .supportCacheLoad()
    //   .then((falg) => falg && this.loadHander.load())
  }
  public getLoadHander(): BaseLoadHandler {
    return this.loadHander
  }
  private model?: tf.LayersModel
  public getInputPadding(): number {
    return 4
  }
  public getOutPadding(): number {
    return 16
  }

  public predictSize(width: number, height: number) {
    width *= 4
    height *= 4
    // width += 8
    // height += 8
    return { width, height }
  }
  // public async load(): Promise<void> {
  //   this.model = await loadModel()
  //   super.load()
  // }
  public predict(input: tf.Tensor3D): tf.Tensor3D {
    if (!this.isLoaded) throw new Error('model not Loading')
    const x = input.reshape([1, ...input.shape])
    const y = this.model!.predict(x) as tf.Tensor4D
    const output = y.as3D(y.shape[1], y.shape[2], y.shape[3])
    // console.log(output)
    tf.dispose([x, y])
    return output
  }
}

export default new RealESRGANModel()
